# 处理带有依赖关系的序列（时间）的数据
# 自然语言有序列关系
# RNNcell and RNN
# demo of sequence to sequence:
import torch
import torch.nn as nn
import torch.optim as optim

input_size = 4
hidden_size = 4
batch_size = 1
num_layers = 1
seq_len = 5

idx2char = ['e', 'h', 'l', 'o']
x_data = [1, 0, 2, 2, 3]
y_data = [3, 1, 2, 3, 2]
one_hot_look_up = [
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 1, 0],
    [0, 0, 0, 1]
]
x_one_hot = [one_hot_look_up[x] for x in x_data]
inputs = torch.Tensor(x_one_hot).view(seq_len, batch_size, input_size)
labels = torch.LongTensor(y_data)


class rnn_Model(nn.Module):
    def __init__(self, input_size, hidden_size, batch_size, num_layers=1):
        super(rnn_Model, self).__init__()
        self.num_layer = num_layers
        self.batch_size = batch_size
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.rnn = torch.nn.RNN(input_size=self.input_size, hidden_size=self.hidden_size, num_layers=self.num_layer)

    def forward(self, input):
        hidden = torch.zeros(self.num_layer, self.batch_size, self.hidden_size)
        out, _ = self.rnn(input, hidden)
        return out.view(-1, self.hidden_size)


net = rnn_Model(input_size, hidden_size, batch_size, num_layers)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.01)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net.to(device)
print(f'使用设备:{device}')

for epoch in range(100):
    optimizer.zero_grad()
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

    _, idx = outputs.max(dim=1)
    idx = idx.data.numpy()
    print('predict:', ''.join([idx2char[x] for x in idx]), end='')
    print(',epoch [%d/20] loss = %3f' % (epoch + 1, loss.item()))


# embedding:数据降维